import unittest

from tinydb.sql.optimizier.planner import query_logical_plan, query_physical_plan

from tinydb.sql.parser.parser import query_parse

def explain(operator, indent = ''):
    output_lines = []
    name = operator.name

    if indent:
        output_lines.append(f'{indent} -> {name}')
    else:
        # the first line
        output_lines.append(f'{name}')

    for i, child in enumerate(operator.children):
        child_indent = f'{indent}  '
        output_lines.extend(explain(child, child_indent))
    return output_lines

class TestPlanner(unittest.TestCase):
    def test_logical_plan(self):
        ast = query_parse('select * from t1 where t1.id > 0')
        query = query_logical_plan(ast)

        ret = '\n'.join(explain(query))
        self.assertEqual(ret, 'Query\n   -> Scan')

        physical_plan = query_physical_plan(query)
        physical_plan.open()
        result = list(physical_plan.next())
        print(result)
        physical_plan.close()
        print(physical_plan.elapsed_time, physical_plan.actual_rows)
        ret = '\n'.join(explain(physical_plan))
        self.assertEqual(ret, 'Results\n   -> IndexScan')
        expected = [(1, '小明'), (2, '小红'), (3, '小刚'), (4, '小丽')]
        self.assertEqual(result, expected)

        ast = query_parse('select count(t1.id) from t1 group by t1.id')
        query = query_logical_plan(ast)
        ret = '\n'.join(explain(query))
        self.assertEqual(ret, 'Query\n   -> Group\n     -> Scan')

if __name__ == '__main__':
    unittest.main()